feat(server): implement Resource Scoping for tasks and push notifications#709
feat(server): implement Resource Scoping for tasks and push notifications#709sokoliva wants to merge 42 commits intoa2aproject:1.0-devfrom
Resource Scoping for tasks and push notifications#709Conversation
…fications` Introduces caller indentity isolation to ensure clients only access authorized resources, as mandated by the A2A spec. - Add 'owner' field to `TaskMixin` and `PushNotificationConfig` database models. - Add 'last_updated' field to `TaskMixin` for optimized sorting and indexing. - Update `DatabaseTaskStore`, `InMemoryTaskStore` and `DatabasePushNotificationConfigStore` to use `OwnerResolver`. - Add relevant Unit tests. - Add Alembic configuration to enable users to update their own databases with non-optional `owner` field in `tasks` table.
Summary of ChangesHello @sokoliva, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the server's data management capabilities by implementing robust resource scoping for tasks and push notification configurations, ensuring that users can only access their own data. It also introduces a flexible 'ListTasks' API, enabling efficient retrieval and management of tasks with advanced filtering and pagination. The integration of Alembic provides a structured approach to future database schema evolution. Highlights
Changelog
Ignored Files
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces resource scoping by adding an owner field to tasks and push notification configurations, which is a great step towards multi-tenancy and security. The implementation across the database and in-memory stores is consistent. I've also noticed the addition of Alembic for database migrations, which is excellent for managing schema changes. My feedback focuses on improving the Alembic configuration, strengthening the data model for timestamps, and making the owner resolution logic more robust.
…to resource-scoping
- fix elmbic README.md error - make ServerCallContext optional in OwnerResolver
…thon into resource-scoping
…to resource-scoping
…thon into resource-scoping
src/a2a/migrations/README.md
Outdated
|
|
||
| ```bash | ||
| # Bring the database to the latest version | ||
| a2a-db |
There was a problem hiding this comment.
Running this command in a fresh environment will fail with alembic package missing error. alembic is a dev dependency currently so for package users it won't be installed.
You can reproduce it by creating a "fake" consumer project and installing whl there:
uv build
which is going to output something like dist/a2a_python-*.whl.
And in the new folder
uv venv
uv pip install <whl path from uv build>
uv run a2a-db
I believe we have two options:
- Instruct users to manually install
alembicas their dev dependency if they want to use this CLI. - Create a separate "extra" which will be used for this.
There was a problem hiding this comment.
Added a new optional dependency cli = ["alembic>=1.14.0"] and ImportError handling to cli.py, env.py and 6419d2d130f6_add_columns_owner_last_updated.py. WDYT?
There was a problem hiding this comment.
SG, let's update the readme as well.
| self.owner_resolver = owner_resolver | ||
|
|
||
| def _get_owner_tasks(self, owner: str) -> dict[str, Task]: | ||
| return self.tasks.get(owner, {}) |
There was a problem hiding this comment.
I don't think that fallback to {} is required after changing tasks to defaultdict(dict).
There was a problem hiding this comment.
Without fallback to { } owner_tasks = self._get_owner_tasks(owner) become of type dict[str, Task] | None which makes task = owner_tasks.get(task_id) fail with "AttributeError: 'NoneType' object has no attribute 'get'".
There was a problem hiding this comment.
Is it a runtime or a type error? Accessing an unknown key for defaultdict should produce dict. .get(Key) returns dict[str, Task] | None but for defaultdict it's safe to use indexer:
| return self.tasks.get(owner, {}) | |
| return self.tasks[owner] |
… and last_updated columns to tasks and push_notification_configs
…thon into resource-scoping
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces resource scoping based on an owner field, which is a significant and well-implemented feature. The changes span across database models, data stores (both database and in-memory), and request handlers, ensuring consistent data isolation. The inclusion of Alembic for database migrations, complete with a CLI tool and thorough documentation, is a great addition that will make it much easier for users to manage database schema updates. The new tests, especially the ones for resource scoping and the migration scripts, are comprehensive and add a lot of value. I have a few minor suggestions to improve maintainability and robustness.
src/a2a/migrations/versions/6419d2d130f6_add_columns_owner_last_updated.py
Show resolved
Hide resolved
| tables_str = context.config.get_main_option( | ||
| 'tables', 'tasks,push_notification_configs' | ||
| ) |
There was a problem hiding this comment.
Let's have two explicit different flags, one for each table and handle each table explicitly even if it introduces some code duplication.
| #### -o | ||
| Allows you to pass custom values for the new `owner` column. If not set, it will default to the value `unknown`. |
There was a problem hiding this comment.
Consider renaming to something more specific without short option. Going forward with more migrations coming we may want to introduce other parameters for other migrations.
One option is to prefix it with some disambiguated migration name.
| uv run a2a-db -v | ||
| ``` | ||
| #### --sql | ||
| Enables running migrations in `offline` mode. Migrations are generated as SQL scripts instead of being run against the database. |
There was a problem hiding this comment.
Let's specify exactly how "Migrations are generated as SQL scripts" - printed to stdout or some file.
| if not context: | ||
| return 'unknown' | ||
| # Example: Basic user name. Adapt as needed for your user model. | ||
| return context.user.user_name |
There was a problem hiding this comment.
We need to be careful with defaults here and how those defaults interact with DB migration defaults.
Using unknown in both places will lead to allowing everyone to access prior data.
Maybe DB default should be something more specific like legacy_v03_no_user_info.
| def resolve_user_scope(context: ServerCallContext | None) -> str: | ||
| """Resolves the owner scope based on the user in the context.""" | ||
| if not context: | ||
| return 'unknown' |
There was a problem hiding this comment.
In which cases we won't have context here?
| mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"] | ||
| signing = ["PyJWT>=2.0.0"] | ||
| sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"] | ||
| cli = ["alembic>=1.14.0"] |
There was a problem hiding this comment.
Let's consolidate naming with db part, the CLI is named a2a-db as a command, while it is in cli.py and also has just cli here.
To give us more room in the future I suggest to use db everywhere: extra name, file name, command name.
| self.owner_resolver = owner_resolver | ||
|
|
||
| def _get_owner_tasks(self, owner: str) -> dict[str, Task]: | ||
| return self.tasks.get(owner, {}) |
There was a problem hiding this comment.
Is it a runtime or a type error? Accessing an unknown key for defaultdict should produce dict. .get(Key) returns dict[str, Task] | None but for defaultdict it's safe to use indexer:
| return self.tasks.get(owner, {}) | |
| return self.tasks[owner] |
src/a2a/migrations/README.md
Outdated
|
|
||
| ```bash | ||
| # Bring the database to the latest version | ||
| a2a-db |
There was a problem hiding this comment.
SG, let's update the readme as well.
3 participants
Description
Introduces caller indentity isolation to ensure clients only access authorized resources, as mandated by the A2A spec.
ownerfield toTaskMixinandPushNotificationConfigdatabase models.last_updatedfield toTaskMixinfor optimized sorting and indexing.DatabaseTaskStore,InMemoryTaskStoreandDatabasePushNotificationConfigStoreto useOwnerResolver.ownercolumn intasksandpush_notification_configstable and optionallast_updatedand index(owner, last_updated)intasks.uv run a2a-dbfor database updating.Note
src/a2a/server/tasks/database_task_store.pylistmethod, Gemini suggested a refactor of pagination. I thoroughly reviewed it and comfirmed that the logic is the same and that readability of code improved so I decided to accept it.delete_infomethod. Whenconfig_idis None and onlytask_idwas provided it would search for configs mapped totask_idwithconfig.id=task_id, contrary todelete_infomethod of DatabasePushNotificationConfigStore where if config_id is None, all configurations for the task are deleted. Unfortunately, I did not find intended behavior defined in the spec, but behavior ofDatabasePushNotificationConfigStore'sdelete_infoseems more logical.Breaking changes
added non-optional owner field to the Task Model. Use alembic configuration to update your database.
Ensure the tests and linter pass (Run
bash scripts/format.shfrom the repository root to format)Appropriate docs were updated (if necessary)
Fixes #610 🦕